package cn.xeblog.server.action.timer;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class TokenBucketFactory {
    private static final Map<String, TokenBucket> tokenBuckets = new ConcurrentHashMap<>();

    static {
        ScheduledExecutorService refillExecutor = Executors.newScheduledThreadPool(1);
        refillExecutor.scheduleAtFixedRate(TokenBucketFactory::refillAllTokens, 0, 1, TimeUnit.SECONDS);
    }

    private TokenBucketFactory() {
    }

    public static void registerTokenBucket(TokenBucket tokenBucket) {
        tokenBuckets.put(tokenBucket.getBucketId(), tokenBucket);
    }

    public static void removeTokenBucket(String bucketId) {
        tokenBuckets.remove(bucketId);
    }

    public static TokenBucket getTokenBucket(String bucketId) {
        return tokenBuckets.get(bucketId);
    }

    private static void refillAllTokens() {
        tokenBuckets.values().forEach(TokenBucket::refillTokens);
    }
}
